% function Rxxb = mac2bc_Lxu_v2(Rxxm, Hmac, Lxu)
%
% MAC-to-BC covariance duality conversion, allowing different numbers of
% transmit antennas for different MAC users.
%
% Backward-compatible uniform-user use:
%     Rxxb = mac2bc_Lxu_v2(Rxxm, Hmac)
% where Hmac is Ly x Lxu x U and Rxxm is Lxu x Lxu x U.
%
% Unequal-user use:
%     Rxxb = mac2bc_Lxu_v2(Rxxm, Hmac, Lxu)
% where Lxu = [Lx1 ... LxU].
%
% Accepted Hmac formats for unequal Lxu:
%   (1) Ly x sum(Lxu) x N          users concatenated in columns
%   (2) Ly x 1 x sum(Lxu) x N      one column per antenna/stream
%   (3) Ly x sum(Lxu)              one tone, users concatenated
%   (4) Ly x 1 x sum(Lxu)          one tone, one column per antenna/stream
%
% Accepted Rxxm formats:
%   (1) U x N cell array, Rxxm{u,n} is Lxu(u) x Lxu(u)
%   (2) Lmax x Lmax x U x N padded numeric array
%   (3) sum(Lxu) x sum(Lxu) x N block-diagonal numeric array
%   (4) sum(Lxu) x sum(Lxu) one-tone block-diagonal numeric array
%
% Output:
%   Rxxb is Ly x Ly x U x N.  If N=1, output is Ly x Ly x U.
%
% The input users are internally reversed exactly as in the original
% mac2bc.m code, so that the loop index k follows the textbook ordering.

function Rxxb = mac2bc_Lxu_v2(Rxxm, Hmac, Lxu)

    if nargin < 3 || isempty(Lxu)
        % Original uniform-input convention: Hmac is Ly x Lx_per_user x U.
        [Ly, Lper, U] = size(Hmac);
        N = 1;
        Lxu = Lper*ones(1,U);

        Hcell = cell(U,N);
        Rcell = cell(U,N);
        for u = 1:U
            Hcell{u,1} = Hmac(:,:,u);
            Rcell{u,1} = Rxxm(:,:,u);
        end
    else
        Lxu = reshape(Lxu,1,[]);
        U = numel(Lxu);
        Ltot_user = sum(Lxu);

        [Hcat, Ly, Ltot, N] = normalize_Hmac(Hmac, Ltot_user);
        if Ltot ~= Ltot_user
            error('sum(Lxu) = %d must equal total MAC transmit columns = %d.', Ltot_user, Ltot);
        end

        idx_end   = cumsum(Lxu);
        idx_start = [1, idx_end(1:end-1)+1];

        Hcell = cell(U,N);
        for n = 1:N
            for u = 1:U
                Hcell{u,n} = Hcat(:,idx_start(u):idx_end(u),n);
            end
        end

        Rcell = unpack_Rxxm(Rxxm, Lxu, idx_start, idx_end, U, N);
    end

    % Reverse users to preserve the convention in the original mac2bc.m.
    rev = U:-1:1;
    Hcell = Hcell(rev,:);
    Rcell = Rcell(rev,:);
    Lxu   = Lxu(rev);

    Rxxb = zeros(Ly,Ly,U,N);

    for n = 1:N
        B = cell(U,1);
        B{U} = eye(Ly);
        for k = U:-1:2
            Hk = Hcell{k,n};
            Rk = hermitianize(Rcell{k,n});
            B{k-1} = hermitianize(B{k} + Hk*Rk*Hk');
        end

        Rxxbtot = zeros(Ly,Ly);

        for k = 1:U
            Hk = Hcell{k,n};
            Rk = hermitianize(Rcell{k,n});
            Lk = Lxu(k);

            if k == 1
                Ak = eye(Lk);
            else
                Ak = hermitianize(eye(Lk) + Hk'*Rxxbtot*Hk);
            end

            invsqrt_A = invsqrtm_hpd(Ak);
            invsqrt_B = invsqrtm_hpd(B{k});
            sqrt_A    = sqrtm_hpd(Ak);

            [F,~,M] = svd(invsqrt_B' * Hk * invsqrt_A', 'econ');

            Rb = invsqrt_B' * F * M' * sqrt_A * Rk * sqrt_A' * M * F' * invsqrt_B;
            Rb = hermitianize(Rb);

            Rxxb(:,:,k,n) = Rb;
            Rxxbtot = hermitianize(Rxxbtot + Rb);
        end
    end

    if N == 1
        Rxxb = Rxxb(:,:,:,1);
    end
end

function [Hcat, Ly, Ltot, N] = normalize_Hmac(Hmac, Ltot_expected)
    sz = size(Hmac);
    nd = ndims(Hmac);
    Ly = sz(1);

    if ismatrix(Hmac)
        % Ly x Ltot, one tone.
        Ltot = sz(2);
        N = 1;
        Hcat = reshape(Hmac, Ly, Ltot, 1);
        return;
    end

    if nd == 3
        if sz(2) == Ltot_expected
            % Ly x Ltot x N.
            Ltot = sz(2);
            N = sz(3);
            Hcat = Hmac;
        elseif sz(2) == 1 && sz(3) == Ltot_expected
            % Ly x 1 x Ltot, one tone with columns stored in third dim.
            Ltot = sz(3);
            N = 1;
            Hcat = reshape(Hmac, Ly, Ltot, 1);
        else
            error('Unrecognized 3-D Hmac format. Use Ly x sum(Lxu) x N or Ly x 1 x sum(Lxu).');
        end
        return;
    end

    if nd == 4
        if sz(2) == 1 && sz(3) == Ltot_expected
            % Ly x 1 x Ltot x N.
            Ltot = sz(3);
            N = sz(4);
            Hcat = zeros(Ly,Ltot,N);
            for n = 1:N
                Hcat(:,:,n) = reshape(Hmac(:,:,:,n), Ly, Ltot);
            end
        else
            error('Unrecognized 4-D Hmac format. Expected Ly x 1 x sum(Lxu) x N.');
        end
        return;
    end

    error('Unrecognized Hmac format.');
end

function Rcell = unpack_Rxxm(Rxxm, Lxu, idx_start, idx_end, U, N)
    Rcell = cell(U,N);

    if iscell(Rxxm)
        if isequal(size(Rxxm), [U N])
            for n = 1:N
                for u = 1:U
                    Rcell{u,n} = Rxxm{u,n};
                end
            end
        elseif N == 1 && numel(Rxxm) == U
            for u = 1:U
                Rcell{u,1} = Rxxm{u};
            end
        else
            error('Cell Rxxm must be U x N, or U x 1 when N=1.');
        end
        return;
    end

    sz = size(Rxxm);

    % Padded format: Lmax x Lmax x U x N.
    if numel(sz) >= 3 && sz(3) == U
        if numel(sz) < 4
            sz(4) = 1;
        end
        if sz(4) ~= N
            error('Rxxm has %d tones but Hmac has %d tones.', sz(4), N);
        end
        for n = 1:N
            for u = 1:U
                Lu = Lxu(u);
                Rcell{u,n} = Rxxm(1:Lu,1:Lu,u,n);
            end
        end
        return;
    end

    % Block-diagonal format: sum(Lxu) x sum(Lxu) x N, or one-tone 2-D.
    if sz(1) == sum(Lxu) && sz(2) == sum(Lxu)
        if numel(sz) < 3
            Nt_R = 1;
        else
            Nt_R = sz(3);
        end
        if Nt_R ~= N
            error('Block Rxxm has %d tones but Hmac has %d tones.', Nt_R, N);
        end
        for n = 1:N
            for u = 1:U
                ii = idx_start(u):idx_end(u);
                if ndims(Rxxm) < 3
                    Rcell{u,n} = Rxxm(ii,ii);
                else
                    Rcell{u,n} = Rxxm(ii,ii,n);
                end
            end
        end
        return;
    end

    error('Unrecognized Rxxm format. Use cell, padded, or block-diagonal format.');
end

function X = hermitianize(X)
    X = 0.5*(X + X');
end

function S = sqrtm_hpd(A)
    A = hermitianize(A);
    [V,D] = eig(A);
    d = max(real(diag(D)), 0);
    S = V*diag(sqrt(d))*V';
    S = hermitianize(S);
end

function S = invsqrtm_hpd(A)
    A = hermitianize(A);
    [V,D] = eig(A);
    d = max(real(diag(D)), 1e-14);
    S = V*diag(1./sqrt(d))*V';
    S = hermitianize(S);
end
